#!/usr/bin/env python3
"""
PLOOP数据处理器（增强版） - 基于IDA Pro分析的完整SET2D算法实现
解决关键问题：
1. 实现空间索引系统（对应D2_LimitsArray）
2. 边拓扑管理（对应平行数组结构）
3. 改进数值稳定性
4. 严格的路径验证和优化
"""

import math
import re
from typing import List, Tuple, Optional, Dict, Set
from dataclasses import dataclass, field
from collections import defaultdict
import json

@dataclass
class Vertex:
    """增强的顶点数据结构"""
    x: float
    y: float
    z: float
    fradius: Optional[float] = None
    vertex_id: Optional[int] = None  # 顶点唯一标识
    
    def __str__(self):
        if self.fradius:
            return f"({self.x:.2f}, {self.y:.2f}, {self.z:.2f}, FRAD:{self.fradius:.2f})"
        return f"({self.x:.2f}, {self.y:.2f}, {self.z:.2f})"
    
    def distance_to(self, other: 'Vertex') -> float:
        """计算到另一个顶点的距离"""
        return math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
    
    def equals(self, other: 'Vertex', tolerance: float = 0.001) -> bool:
        """检查是否与另一个顶点相等（在容差范围内）"""
        return self.distance_to(other) <= tolerance

@dataclass
class Edge:
    """边数据结构（对应IDA Pro的边索引系统）"""
    start_vertex_id: int
    end_vertex_id: int
    label: int = 0
    edge_id: Optional[int] = None
    is_arc: bool = False  # 是否为圆弧边
    arc_data: Optional[dict] = None  # 圆弧数据
    
    def is_degenerate(self) -> bool:
        """检查是否为退化边（起点和终点相同）"""
        return self.start_vertex_id == self.end_vertex_id

class SpatialIndex:
    """空间索引系统（对应IDA Pro的D2_LimitsArray）"""
    
    def __init__(self, tolerance: float = 0.1):
        self.tolerance = tolerance
        self.vertex_grid: Dict[Tuple[int, int], List[int]] = defaultdict(list)
        self.grid_size = max(tolerance * 10, 1.0)  # 网格大小
        self.vertices: Dict[int, Vertex] = {}
        
    def _get_grid_key(self, x: float, y: float) -> Tuple[int, int]:
        """获取网格键值"""
        return (int(x // self.grid_size), int(y // self.grid_size))
    
    def add_vertex(self, vertex: Vertex, vertex_id: int):
        """添加顶点到空间索引"""
        self.vertices[vertex_id] = vertex
        grid_key = self._get_grid_key(vertex.x, vertex.y)
        self.vertex_grid[grid_key].append(vertex_id)
    
    def find_nearby_vertices(self, x: float, y: float, radius: float = None) -> List[int]:
        """查找附近的顶点（对应D2_LimitsArrayIterator的功能）"""
        if radius is None:
            radius = self.tolerance
            
        nearby_vertices = []
        grid_key = self._get_grid_key(x, y)
        
        # 搜索周围的网格
        grid_radius = int(radius // self.grid_size) + 1
        for dx in range(-grid_radius, grid_radius + 1):
            for dy in range(-grid_radius, grid_radius + 1):
                search_key = (grid_key[0] + dx, grid_key[1] + dy)
                for vertex_id in self.vertex_grid.get(search_key, []):
                    vertex = self.vertices[vertex_id]
                    distance = math.sqrt((vertex.x - x)**2 + (vertex.y - y)**2)
                    if distance <= radius:
                        nearby_vertices.append(vertex_id)
        
        return nearby_vertices
    
    def clear(self):
        """清空空间索引"""
        self.vertex_grid.clear()
        self.vertices.clear()

class EdgeTopology:
    """边拓扑管理系统（对应IDA Pro的平行数组结构）"""
    
    def __init__(self):
        self.edges: List[Edge] = []
        self.vertex_to_edges: Dict[int, List[int]] = defaultdict(list)
        self.next_edge_id = 0
        
    def add_edge(self, start_vertex_id: int, end_vertex_id: int, label: int = 0) -> int:
        """添加边"""
        edge = Edge(
            start_vertex_id=start_vertex_id,
            end_vertex_id=end_vertex_id,
            label=label,
            edge_id=self.next_edge_id
        )
        
        edge_id = self.next_edge_id
        self.edges.append(edge)
        self.vertex_to_edges[start_vertex_id].append(edge_id)
        self.vertex_to_edges[end_vertex_id].append(edge_id)
        self.next_edge_id += 1
        
        return edge_id
    
    def remove_edge(self, edge_id: int):
        """移除边"""
        if edge_id < len(self.edges):
            edge = self.edges[edge_id]
            edge.start_vertex_id = -1  # 标记为删除
            edge.end_vertex_id = -1

class EnhancedPLoopProcessor:
    """增强版PLOOP处理器"""
    
    def __init__(self, tolerance: float = 0.1):
        self.tolerance = tolerance
        self.spatial_index = SpatialIndex(tolerance)
        self.edge_topology = EdgeTopology()
        self.vertices: List[Vertex] = []
        self.vertex_id_counter = 0
        
        # 数值稳定性参数
        self.min_arc_angle = 0.001  # 最小圆弧角度
        self.max_radius_ratio = 1000.0  # 最大半径比例
        
    def parse_position(self, pos_str: str) -> Tuple[float, float, float]:
        """解析位置字符串"""
        x, y, z = 0.0, 0.0, 0.0
        
        # 解析E/W方向（X轴）
        e_match = re.search(r'E\s+([\d.-]+)mm', pos_str)
        w_match = re.search(r'W\s+([\d.-]+)mm', pos_str)
        if e_match:
            x = float(e_match.group(1))
        elif w_match:
            x = -float(w_match.group(1))
        
        # 解析N/S方向（Y轴）
        n_match = re.search(r'N\s+([\d.-]+)mm', pos_str)
        s_match = re.search(r'S\s+([\d.-]+)mm', pos_str)
        if n_match:
            y = float(n_match.group(1))
        elif s_match:
            y = -float(s_match.group(1))
        
        # 解析U方向（Z轴）
        u_match = re.search(r'U\s+([\d.-]+)mm', pos_str)
        if u_match:
            z = float(u_match.group(1))
        
        return x, y, z
    
    def parse_file(self, content: str) -> List['PLoop']:
        """解析PLOOP文件"""
        ploops = []
        lines = content.split('\n')
        
        framework_name = "Unknown"
        current_ploop = None
        current_vertex = None
        end_count = 0
        
        for line_num, line in enumerate(lines, 1):
            line = line.strip()
            
            if line.startswith('NEW FRMWORK'):
                parts = line.split()
                framework_name = parts[2] if len(parts) > 2 else f"Framework_{line_num}"
                print(f"发现FRMWORK: {framework_name}")
            
            elif line == 'NEW PLOOP':
                current_ploop = {'name': framework_name, 'vertices': [], 'height': 0.0}
                end_count = 0
                print(f"  开始PLOOP: {framework_name}")
            
            elif line.startswith('HEIG') and current_ploop is not None:
                height_match = re.search(r'HEIG\s+([\d.-]+)mm', line)
                if height_match:
                    current_ploop['height'] = float(height_match.group(1))
                    print(f"    高度: {current_ploop['height']}mm")
            
            elif line == 'NEW PAVERT' and current_ploop is not None:
                current_vertex = {'x': 0.0, 'y': 0.0, 'z': 0.0, 'fradius': None}
            
            elif line.startswith('POS') and current_vertex is not None:
                x, y, z = self.parse_position(line)
                current_vertex.update({'x': x, 'y': y, 'z': z})
            
            elif line.startswith('FRAD') and current_vertex is not None:
                frad_match = re.search(r'FRAD\s+([\d.-]+)mm', line)
                if frad_match:
                    current_vertex['fradius'] = float(frad_match.group(1))
                    print(f"      发现FRAD: {current_vertex['fradius']}mm")
            
            elif line == 'END':
                if current_vertex is not None:
                    vertex = Vertex(
                        x=current_vertex['x'],
                        y=current_vertex['y'], 
                        z=current_vertex['z'],
                        fradius=current_vertex['fradius'],
                        vertex_id=len(current_ploop['vertices'])
                    )
                    current_ploop['vertices'].append(vertex)
                    print(f"      添加顶点: {vertex}")
                    current_vertex = None
                elif current_ploop is not None:
                    end_count += 1
                    if end_count >= 3:
                        if len(current_ploop['vertices']) > 0:
                            ploop = PLoop(
                                name=current_ploop['name'],
                                height=current_ploop['height'],
                                vertices=current_ploop['vertices']
                            )
                            ploops.append(ploop)
                            print(f"  完成PLOOP: {ploop}")
                        current_ploop = None
                        end_count = 0
        
        return ploops
    
    def add_vertex_with_merge(self, vertex: Vertex) -> int:
        """添加顶点，如果附近已有顶点则合并（对应IDA Pro的adjustVertices）"""
        # 查找附近的顶点
        nearby_vertices = self.spatial_index.find_nearby_vertices(
            vertex.x, vertex.y, self.tolerance
        )
        
        # 检查是否可以合并
        for vertex_id in nearby_vertices:
            existing_vertex = self.vertices[vertex_id]
            if vertex.equals(existing_vertex, self.tolerance):
                print(f"    合并顶点: ({vertex.x:.2f}, {vertex.y:.2f}) → 顶点[{vertex_id}]")
                # 保留FRADIUS信息（如果新顶点有的话）
                if vertex.fradius and not existing_vertex.fradius:
                    existing_vertex.fradius = vertex.fradius
                return vertex_id
        
        # 添加新顶点
        vertex_id = self.vertex_id_counter
        vertex.vertex_id = vertex_id
        self.vertices.append(vertex)
        self.spatial_index.add_vertex(vertex, vertex_id)
        self.vertex_id_counter += 1
        
        return vertex_id
    
    def calculate_robust_fillet_arc(self, prev_vertex: Vertex, vertex: Vertex, next_vertex: Vertex) -> Optional[dict]:
        """改进的圆弧计算，具有更好的数值稳定性"""
        if not vertex.fradius or vertex.fradius <= 0:
            return None
        
        # 计算边向量
        v1_x = vertex.x - prev_vertex.x
        v1_y = vertex.y - prev_vertex.y
        v2_x = next_vertex.x - vertex.x
        v2_y = next_vertex.y - vertex.y
        
        # 边长度
        len1 = math.sqrt(v1_x*v1_x + v1_y*v1_y)
        len2 = math.sqrt(v2_x*v2_x + v2_y*v2_y)
        
        # 数值稳定性检查
        if len1 < self.tolerance or len2 < self.tolerance:
            print(f"    警告: 边长度过小，跳过圆角计算")
            return None
        
        # 标准化向量
        v1_x /= len1
        v1_y /= len1
        v2_x /= len2
        v2_y /= len2
        
        # 计算夹角
        dot_product = v1_x * v2_x + v1_y * v2_y
        dot_product = max(-1.0, min(1.0, dot_product))
        
        angle = math.acos(abs(dot_product))
        half_angle = angle / 2.0
        
        # 角度有效性检查
        if half_angle < self.min_arc_angle:
            print(f"    警告: 角度过小 ({math.degrees(angle):.2f}°)，跳过圆角")
            return None
        
        # 计算到切点的距离
        tan_half = math.tan(half_angle)
        if tan_half < self.tolerance:
            return None
        
        dist_to_tangent = vertex.fradius / tan_half
        
        # 半径合理性检查
        max_dist = min(len1, len2) * 0.9
        if dist_to_tangent > max_dist:
            # 调整半径以适应几何约束
            adjusted_radius = max_dist * tan_half
            print(f"    调整圆角半径: {vertex.fradius:.2f} → {adjusted_radius:.2f}mm")
            dist_to_tangent = max_dist
            vertex.fradius = adjusted_radius
        
        # 计算切点
        tangent1_x = vertex.x - v1_x * dist_to_tangent
        tangent1_y = vertex.y - v1_y * dist_to_tangent
        tangent2_x = vertex.x + v2_x * dist_to_tangent
        tangent2_y = vertex.y + v2_y * dist_to_tangent
        
        # 计算圆心
        bisector_x = v1_x + v2_x
        bisector_y = v1_y + v2_y
        bisector_len = math.sqrt(bisector_x*bisector_x + bisector_y*bisector_y)
        
        if bisector_len < self.tolerance:
            return None
        
        bisector_x /= bisector_len
        bisector_y /= bisector_len
        
        # 判断凸凹性
        cross_product = v1_x * v2_y - v1_y * v2_x
        radius_sign = 1.0 if cross_product > 0 else -1.0
        
        center_dist = vertex.fradius / math.sin(half_angle)
        center_x = vertex.x + bisector_x * center_dist * radius_sign
        center_y = vertex.y + bisector_y * center_dist * radius_sign
        
        # 计算角度信息
        angle1 = math.atan2(tangent1_y - center_y, tangent1_x - center_x)
        angle2 = math.atan2(tangent2_y - center_y, tangent2_x - center_x)
        
        angle_diff = angle2 - angle1
        if abs(angle_diff) > math.pi:
            if angle_diff > 0:
                angle_diff -= 2 * math.pi
            else:
                angle_diff += 2 * math.pi
        
        return {
            'center': (center_x, center_y),
            'radius': vertex.fradius,
            'start_point': (tangent1_x, tangent1_y),
            'end_point': (tangent2_x, tangent2_y),
            'start_angle': angle1,
            'end_angle': angle2,
            'angle_diff': angle_diff,
            'large_arc': abs(angle_diff) > math.pi,
            'sweep_flag': angle_diff > 0,
            'original_vertex': vertex
        }
    
    def process_ploop_enhanced(self, ploop: 'PLoop') -> List[Vertex]:
        """增强的PLOOP处理"""
        if len(ploop.vertices) < 3:
            return ploop.vertices
        
        print(f"\n🔧 增强处理PLOOP: {ploop.name}")
        print(f"  原始顶点数: {len(ploop.vertices)}")
        
        # 1. 重置处理器状态
        self.spatial_index.clear()
        self.edge_topology = EdgeTopology()
        self.vertices.clear()
        self.vertex_id_counter = 0
        
        # 2. 添加所有顶点（自动合并重复顶点）
        vertex_ids = []
        for i, vertex in enumerate(ploop.vertices):
            vertex_id = self.add_vertex_with_merge(vertex)
            vertex_ids.append(vertex_id)
        
        print(f"  合并后顶点数: {len(self.vertices)}")
        
        # 3. 建立边拓扑
        for i in range(len(vertex_ids)):
            start_id = vertex_ids[i]
            end_id = vertex_ids[(i + 1) % len(vertex_ids)]
            if start_id != end_id:  # 避免自环
                self.edge_topology.add_edge(start_id, end_id)
        
        # 4. 处理FRADIUS圆角
        processed_vertices = []
        n = len(self.vertices)
        
        for i in range(n):
            current_vertex = self.vertices[i]
            
            # 查找前后顶点
            prev_vertex = self.vertices[i - 1] if i > 0 else self.vertices[n - 1]
            next_vertex = self.vertices[i + 1] if i < n - 1 else self.vertices[0]
            
            if current_vertex.fradius and current_vertex.fradius > 0:
                # 计算圆弧信息
                arc_info = self.calculate_robust_fillet_arc(prev_vertex, current_vertex, next_vertex)
                if arc_info:
                    # 添加圆弧起点（带FRADIUS信息用于SVG绘制）
                    start_vertex = Vertex(
                        arc_info['start_point'][0], 
                        arc_info['start_point'][1], 
                        current_vertex.z,
                        current_vertex.fradius
                    )
                    # 添加圆弧终点
                    end_vertex = Vertex(
                        arc_info['end_point'][0], 
                        arc_info['end_point'][1], 
                        current_vertex.z
                    )
                    
                    processed_vertices.extend([start_vertex, end_vertex])
                    print(f"    处理圆角[{i}]: 半径{current_vertex.fradius:.2f}mm")
                else:
                    processed_vertices.append(current_vertex)
            else:
                processed_vertices.append(current_vertex)
        
        # 5. 最终路径验证和清理
        final_vertices = self.validate_and_clean_path(processed_vertices)
        
        print(f"  最终顶点数: {len(final_vertices)}")
        return final_vertices
    
    def validate_and_clean_path(self, vertices: List[Vertex]) -> List[Vertex]:
        """验证和清理路径，确保没有重复点和多余点"""
        if len(vertices) < 3:
            return vertices
        
        print("  🔍 路径验证和清理...")
        
        # 1. 移除重复的相邻顶点
        cleaned = []
        for i, vertex in enumerate(vertices):
            if i == 0 or not vertex.equals(vertices[i-1], self.tolerance):
                cleaned.append(vertex)
        
        print(f"    移除重复相邻顶点: {len(vertices)} → {len(cleaned)}")
        
        # 2. 移除共线的中间顶点
        if len(cleaned) >= 3:
            final_cleaned = []
            n = len(cleaned)
            
            for i in range(n):
                current = cleaned[i]
                prev_vertex = cleaned[i - 1] if i > 0 else cleaned[n - 1]
                next_vertex = cleaned[i + 1] if i < n - 1 else cleaned[0]
                
                # 检查是否共线（但保留有FRADIUS的顶点）
                if current.fradius or not self.is_collinear_strict(prev_vertex, current, next_vertex):
                    final_cleaned.append(current)
                else:
                    print(f"    移除共线顶点: ({current.x:.2f}, {current.y:.2f})")
            
            cleaned = final_cleaned
            print(f"    移除共线顶点: {n} → {len(cleaned)}")
        
        # 3. 确保路径闭合
        if len(cleaned) > 2:
            first = cleaned[0]
            last = cleaned[-1]
            if not first.equals(last, self.tolerance):
                # 如果没有闭合，添加闭合点
                closing_vertex = Vertex(first.x, first.y, first.z)
                cleaned.append(closing_vertex)
                print(f"    添加闭合顶点: ({first.x:.2f}, {first.y:.2f})")
            elif first.equals(last, self.tolerance):
                # 如果最后一个点与第一个点重复，移除最后一个点
                cleaned = cleaned[:-1]
                print(f"    移除重复的闭合顶点")
        
        # 4. 最终验证
        self.validate_path_integrity(cleaned)
        
        return cleaned
    
    def is_collinear_strict(self, p1: Vertex, p2: Vertex, p3: Vertex) -> bool:
        """严格的共线检测"""
        v1_x = p2.x - p1.x
        v1_y = p2.y - p1.y
        v2_x = p3.x - p2.x
        v2_y = p3.y - p2.y
        
        cross_product = v1_x * v2_y - v1_y * v2_x
        len1 = math.sqrt(v1_x*v1_x + v1_y*v1_y)
        len2 = math.sqrt(v2_x*v2_x + v2_y*v2_y)
        
        if len1 < self.tolerance or len2 < self.tolerance:
            return False
        
        relative_area = abs(cross_product) / (len1 * len2)
        return relative_area < 0.001  # 严格的共线容差
    
    def validate_path_integrity(self, vertices: List[Vertex]):
        """验证路径完整性"""
        print("  ✅ 路径完整性验证:")
        
        if len(vertices) < 3:
            print("    警告: 顶点数少于3个")
            return
        
        # 检查重复顶点
        duplicate_count = 0
        for i in range(len(vertices)):
            for j in range(i + 1, len(vertices)):
                if vertices[i].equals(vertices[j], self.tolerance):
                    duplicate_count += 1
        
        if duplicate_count > 0:
            print(f"    警告: 发现 {duplicate_count} 对重复顶点")
        else:
            print("    ✓ 无重复顶点")
        
        # 检查边长度
        min_edge_length = float('inf')
        max_edge_length = 0
        for i in range(len(vertices)):
            current = vertices[i]
            next_vertex = vertices[(i + 1) % len(vertices)]
            edge_length = current.distance_to(next_vertex)
            min_edge_length = min(min_edge_length, edge_length)
            max_edge_length = max(max_edge_length, edge_length)
        
        print(f"    边长度范围: {min_edge_length:.2f} ~ {max_edge_length:.2f}mm")
        
        if min_edge_length < self.tolerance:
            print(f"    警告: 存在过短的边 ({min_edge_length:.2f}mm)")
        else:
            print("    ✓ 边长度正常")
        
        # 检查FRADIUS顶点
        fradius_count = sum(1 for v in vertices if v.fradius)
        print(f"    圆角顶点数: {fradius_count}")

@dataclass
class PLoop:
    """PLOOP数据结构"""
    name: str
    height: float
    vertices: List[Vertex]
    
    def __str__(self):
        return f"PLOOP {self.name}: H={self.height}, {len(self.vertices)} vertices"

def test_k717_enhanced():
    """测试K717案例的增强处理"""
    print("🎯 K717增强处理测试")
    print("="*60)
    
    # 创建增强处理器
    processor = EnhancedPLoopProcessor(tolerance=0.1)
    
    # 读取K717数据
    try:
        with open('717.txt', 'r', encoding='utf-8') as f:
            content = f.read()
        print(f"✅ 文件读取成功，内容长度: {len(content)} 字符")
    except FileNotFoundError:
        print("❌ 错误: 找不到文件 717.txt")
        return
    except Exception as e:
        print(f"❌ 读取文件时出错: {e}")
        return
    
    # 解析PLOOP数据
    ploops = processor.parse_file(content)
    print(f"✅ 解析完成，发现 {len(ploops)} 个PLOOP")
    
    # 查找K717
    k717_ploop = None
    for ploop in ploops:
        if 'K717' in ploop.name:
            k717_ploop = ploop
            break
    
    if k717_ploop is None:
        print("❌ 错误: 没有找到K717 PLOOP")
        return
    
    print(f"\n✅ 找到K717: {k717_ploop}")
    
    # 增强处理
    processed_vertices = processor.process_ploop_enhanced(k717_ploop)
    
    # 结果统计
    print(f"\n📊 处理结果统计:")
    print(f"  原始顶点数: {len(k717_ploop.vertices)}")
    print(f"  处理后顶点数: {len(processed_vertices)}")
    print(f"  FRADIUS顶点数: {sum(1 for v in processed_vertices if v.fradius)}")
    
    # 显示处理后的顶点
    print(f"\n📍 处理后的顶点坐标:")
    for i, vertex in enumerate(processed_vertices):
        print(f"  [{i:2d}] {vertex}")
    
    # 生成结果报告
    result = {
        'ploop_name': k717_ploop.name,
        'original_vertex_count': len(k717_ploop.vertices),
        'processed_vertex_count': len(processed_vertices),
        'fradius_count': sum(1 for v in processed_vertices if v.fradius),
        'vertices': [
            {
                'x': v.x,
                'y': v.y,
                'z': v.z,
                'fradius': v.fradius
            } for v in processed_vertices
        ]
    }
    
    # 保存结果
    with open('k717_enhanced_result.json', 'w', encoding='utf-8') as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    print(f"\n✅ 结果已保存到: k717_enhanced_result.json")
    
    return processed_vertices

if __name__ == "__main__":
    test_k717_enhanced() 